package uk.ac.cam.juliet.zedcat.client;

import java.awt.Dimension;
import java.awt.image.BufferedImage;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.concurrent.BlockingQueue;

/**
 *
 * @author David Piggott
 */
public class FrameSubtractAndExtract implements Runnable {

    private final Pipeline mPipeline;
    private final Configuration mConfiguration;
    private final BlockingQueue<Frame> mInputQueue;
    private final BlockingQueue<Frame> mOutputQueue;
    private Dimension mBackgroundDimension;
    private float[] mR = null;
    private float[] mG = null;
    private float[] mB = null;
    private float[] mThresholds = null;
    private int[] mRegionMap = null;

    public FrameSubtractAndExtract(Pipeline p, Configuration c, BlockingQueue<Frame> iQ, BlockingQueue<Frame> oQ) {
        mPipeline = p;
        mConfiguration = c;
        mInputQueue = iQ;
        mOutputQueue = oQ;
    }

    public synchronized void setBackground(BufferedImage bG) {
        int width = bG.getWidth();
        int height = bG.getHeight();
        int[] backgroundPixels = new int[width * height];
        backgroundPixels = bG.getRGB(0, 0, width, height, backgroundPixels, 0, width);
        mR = new float[width * height];
        mG = new float[width * height];
        mB = new float[width * height];
        mThresholds = new float[width * height];
        for (int i = 0; i < width * height; i++) {
            int r = ((backgroundPixels[i] >> 16) & 0x000000FF);
            int g = ((backgroundPixels[i] >> 8) & 0x000000FF);
            int b = ((backgroundPixels[i]) & 0x000000FF);
            mR[i] = r;
            mB[i] = b;
            mG[i] = g;
            mThresholds[i] = 5.0f;
        }
        mBackgroundDimension = new Dimension(width, height);
    }

    public synchronized BufferedImage getBackground() {
        BufferedImage background = null;
        if (mBackgroundDimension != null) {
            int width = mBackgroundDimension.width;
            int height = mBackgroundDimension.height;
            int backgroundPixels[] = new int[width * height];
            for (int i = 0; i < width * height; i++) {
                int r = (((int) mR[i] << 16) & 0x00FF0000);
                int g = (((int) mG[i] << 8) & 0x0000FF00);
                int b = ((int) mB[i] & 0x000000FF);
                backgroundPixels[i] = r | g | b;
            }
            background = new BufferedImage(width, height, BufferedImage.TYPE_INT_BGR);
            background.setRGB(0, 0, width, height, backgroundPixels, 0, width);
        }
        return background;
    }

    public synchronized LinkedList<Region> SubtractAndExtract(Frame frame) {

        int width = frame.getSourceImage().getWidth();
        int height = frame.getSourceImage().getHeight();

        // Check that we have a background and its dimensions
        // are consistent with the current frame
        if (mRegionMap == null
                || mBackgroundDimension == null
                || width != mBackgroundDimension.width
                || height != mBackgroundDimension.height) {
            mPipeline.setBackground();
        }

        // Check region map initialised and of correct size for
        // current frame
        if (mRegionMap == null
                || width * height != mRegionMap.length) {
            mRegionMap = new int[width * height];
        }

        // Subtract background
        LinkedList<Pixel> foregroundList = new LinkedList<Pixel>();
        synchronized (this) {
            for (int i = 0; i < frame.getSourcePixels().length; i++) {
                int r = ((frame.getSourcePixels()[i] >> 16) & 0x000000FF);
                int g = ((frame.getSourcePixels()[i] >> 8) & 0x000000FF);
                int b = ((frame.getSourcePixels()[i]) & 0x000000FF);

                float diff = (r - mR[i]) * (r - mR[i]) + (b - mB[i]) * (b - mB[i]) + (g - mG[i]) * (g - mG[i]);
                if (diff > mConfiguration.getSubtractionThreshold() * mThresholds[i]) {
                    foregroundList.add(new Pixel(i, width));
                }
            }
        }
        frame.setForegroundPixels(foregroundList);

        // Extract regions
        LinkedList<Region> regionList = new LinkedList<Region>();
        if (!mConfiguration.ignoreDisturbance()
                || foregroundList.size()
                < (width * height) * ((double) mConfiguration.maxDisturbance() / (double) 100)) {

            // First we set everything to be unassigned - to do this
            // we choose a number for the cell that won't be used elsewhere
            int emptyMark = width * height + 1;
            for (int i = 0; i < width * height; i++) {
                mRegionMap[i] = emptyMark;
            }

            // We scan left to right, top to bottom through the fg pixels
            for (Pixel pixel : foregroundList) {
                int pX = pixel.getX();
                int pY = pixel.getY();
                int pI = pixel.getArrayIndex();

                boolean leftEmpty = true;
                boolean topEmpty = true;

                // Look left
                if (pX > 0 && mRegionMap[pI - 1] != emptyMark) {
                    leftEmpty = false;
                }

                // Look up
                if (pY > 0 && mRegionMap[pI - width] != emptyMark) {
                    topEmpty = false;
                }

                // This code runs if we might have to merge regions
                if (!topEmpty && !leftEmpty) {
                    // Merge required!
                    // Which two do we merge?
                    int regionLeft = pI - 1;
                    int headLeft = pI - 1;
                    while (regionLeft >= 0) {
                        regionLeft = mRegionMap[regionLeft];
                        if (regionLeft >= 0) {
                            headLeft = regionLeft;
                        }
                    }
                    Region pixelsLeft = regionList.get(-regionLeft - 1);

                    int regionTop = pI - width;
                    int headTop = pI - width;
                    while (regionTop >= 0) {
                        regionTop = mRegionMap[regionTop];
                        if (regionTop >= 0) {
                            headTop = regionTop;
                        }
                    }
                    Region pixelsTop = regionList.get(-regionTop - 1);

                    if (regionTop != regionLeft) {
                        mRegionMap[headLeft] = headTop;
                        pixelsTop.addAll(pixelsLeft);
                        pixelsLeft.clear();
                    } else {
                        mRegionMap[pI] = pI - 1;
                        pixelsLeft.add(pixel);
                    }
                } // This code runs if we have to add the current pixel to the region on its left
                else if (!leftEmpty) {
                    // Add to the region on the left
                    int regionLeft = pI - 1;
                    int headLeft = pI - 1;
                    while (regionLeft >= 0) {
                        regionLeft = mRegionMap[regionLeft];
                        if (regionLeft >= 0) {
                            headLeft = regionLeft;
                        }
                    }
                    Region pixelsLeft = regionList.get(-regionLeft - 1);
                    pixelsLeft.add(pixel);
                    mRegionMap[pI] = pI - 1;
                } // This code runs if we have to add the current pixel to the region above it
                else if (!topEmpty) {
                    // Add to the region on the left
                    int regionTop = pI - width;
                    int headTop = pI - width;
                    while (regionTop >= 0) {
                        regionTop = mRegionMap[regionTop];
                        if (regionTop >= 0) {
                            headTop = regionTop;
                        }
                    }
                    Region pixelsTop = regionList.get(-regionTop - 1);
                    pixelsTop.add(pixel);
                    mRegionMap[pI] = pI - width;
                } // This code runs if it's a new region!
                else {
                    Region region = new Region(frame, mConfiguration);
                    region.add(pixel);
                    regionList.add(region);
                    mRegionMap[pI] = -1 * regionList.size();
                }

            }

            // Remove regions smaller than our tolerance
            Iterator<Region> it = regionList.iterator();
            while (it.hasNext()) {
                Region region = it.next();
                if (region.size() < (width * height) * ((double) mConfiguration.getMinRegionSize() / (double) 100)) {
                    it.remove();
                }
            }

            // Adapt background
            boolean[] adaptationMask = null;
            if (mConfiguration.selectiveAdaptation()) {
                if (mConfiguration.ignoreOnlyRegions()) {
                    adaptationMask = new boolean[width * height];
                    for (Region r : regionList) {
                        for (Pixel p : r) {
                            adaptationMask[p.getArrayIndex()] = true;
                        }
                    }
                } else {
                    adaptationMask = new boolean[width * height];
                    for (Pixel p : foregroundList) {
                        adaptationMask[p.getArrayIndex()] = true;
                    }
                }
            }
            for (int i = 0; i < frame.getSourcePixels().length; i++) {
                int r = ((frame.getSourcePixels()[i] >> 16) & 0x000000FF);
                int g = ((frame.getSourcePixels()[i] >> 8) & 0x000000FF);
                int b = ((frame.getSourcePixels()[i]) & 0x000000FF);
                float diff = (r - mR[i]) * (r - mR[i]) + (b - mB[i]) * (b - mB[i]) + (g - mG[i]) * (g - mG[i]);
                if (!mConfiguration.selectiveAdaptation() || adaptationMask[i] == false) {
                    mR[i] = (1.0f - mConfiguration.getBackgroundAdaptationRate()) * mR[i] + mConfiguration.getBackgroundAdaptationRate() * r;
                    mG[i] = (1.0f - mConfiguration.getBackgroundAdaptationRate()) * mG[i] + mConfiguration.getBackgroundAdaptationRate() * g;
                    mB[i] = (1.0f - mConfiguration.getBackgroundAdaptationRate()) * mB[i] + mConfiguration.getBackgroundAdaptationRate() * b;
                    mThresholds[i] = (1.0f - mConfiguration.getBackgroundAdaptationRate()) * mThresholds[i] + mConfiguration.getBackgroundAdaptationRate() * diff;
                }
            }

        } else {
            mPipeline.appendLine("Significant disturbance detected, re-establishing background.");
            mPipeline.setBackground();
        }
        return regionList;
    }

    public void run() {
        long FPSCountStartTime = System.nanoTime();
        int frameCount = 0;
        while (true) {
            try {
                Frame frame = mInputQueue.take();
                frame.setRegions(SubtractAndExtract(frame));

                long frameCountElapsedTime = System.nanoTime() - FPSCountStartTime;
                if (frameCountElapsedTime > 1000000000L) {
                    FPSCountStartTime = System.nanoTime();
                    mPipeline.updateFramerate(frameCount);
                    frameCount = 1;
                } else {
                    frameCount++;
                }

                mOutputQueue.put(frame);
            } catch (InterruptedException e) {
            }
        }
    }
}
